import numpy as np

# 语种及其对应的数据量
data_list = [
    (375, 'data.list_dutch'),
    (10809, 'data.list_english'),
    (259, 'data.list_french'),
    (26, 'data.list_polish'),
    (221, 'data.list_spanish'),
    (38, 'data.list_portuguese'),
    (470, 'data.list_german'),
    (60, 'data.list_italian')
]

# 上采样因子 alpha
alpha = 0.5

# 总的数据量
N = sum(n for n, _ in data_list)

# 计算每个语种的采样概率
sampling_probs = [(n / N) ** alpha for n, _ in data_list]
sampling_probs = np.array(sampling_probs) / np.sum(sampling_probs)  # 归一化

# 读取所有语种的数据
all_samples = []
for (_, file_path) in data_list:
    with open(file_path, 'r') as f:
        samples = f.readlines()
        all_samples.append(samples)

# 记录每个语种的当前索引
current_indices = [0] * len(data_list)

def get_batch(batch_size):
    batch_samples = []
    global current_indices

    for i, (samples, prob) in enumerate(zip(all_samples, sampling_probs)):
        sample_count = int(batch_size * prob)
        if current_indices[i] + sample_count > len(samples):
            # If not enough samples left, shuffle and start from the beginning
            np.random.shuffle(samples)
            current_indices[i] = 0
        
        batch_samples.extend(samples[current_indices[i]:current_indices[i] + sample_count])
        current_indices[i] += sample_count
    print(current_indices)
    np.random.shuffle(batch_samples)  # Shuffle the batch to mix samples from different languages
    return batch_samples

# Example usage: Get batches of 1000 samples
import pdb; pdb.set_trace()
batch_size = 1000
num_batches = int(np.ceil(N*1.5 / batch_size))

with open('data.list_balance', 'w') as f:
    for _ in range(num_batches):
        batch = get_batch(batch_size)
        f.writelines(batch)

print("Sampling completed and saved to 'data.list_balance'")
